# -*- coding: utf-8 -*-
"""
Created on Fri Aug 21 11:36:39 2020

@author: Farman
"""

import pygrib
import numpy
import joblib
import sys
import pathlib


def update_ma_data(masked_array):
    '''
    Set masked values to zero for further
    '''
    m = masked_array.mask
    d = masked_array.data * (m ^ True)
    masked_array = numpy.ma.masked_array(data=d, mask=m)
    return masked_array


def show_messages(grib_file):
    gribs = pygrib.open(grib_file)
    messages = gribs.read()
    gribs.close()
    
    print('\nMessages in file : %s'%grib_file)
    print('-' * 100)
    
    for m in messages:
        print(m)
    
    print()
    return


def get_grib_data(grib_file, layers):
    gribs = pygrib.open(grib_file)
    messages = gribs.read()
    #print(len(messages))
    gribs.close()
    
    if isinstance(layers, int):
        data = messages[layers].data()
    else:
        data = []
        
        for layer in layers:
            d = messages[layer].data()
            data.append(d)
    
    ma_list = []
    
    for d in data:
        if isinstance(d, numpy.ma.masked_array):
            ma_list.append(update_ma_data(d))
        else:
            ma_list.append(d)
        
    return ma_list


# get_era5_t_e_vsm() and get_era5_t_e_vsm2()
# are two different method of data processing.
# on PC, both of them almost are same in performance.
# use one of them for your convinience.

def get_era5_pe_etc(era5_grib_file):
    '''
    Get 
        1:Potential evaporation:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        2:Runoff:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        3:Surface runoff:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        4:Evaporation:m of water equivalent (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        5:Total precipitation:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
    from file of path 'ERA5_eva_runoff_pre'.
    '''
    gribs = pygrib.open(era5_grib_file)
      
    print('Reading data of hour %2d ...'%0)
    
    df = gribs.read(1)[0].data()[0]
    pe = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    ro = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    sro = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    ev = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    p = df.data * (True ^ df.mask)
    
    for hour in range(1, gribs.messages // 6):
        print('Reading data of hour %2d ...'%hour)
        
        df = gribs.read(1)[0].data()[0]
        pe += df.data * (True ^ df.mask)
    
        df = gribs.read(1)[0].data()[0]
        ro += df.data * (True ^ df.mask)
        
        df = gribs.read(1)[0].data()[0]
        sro += df.data * (True ^ df.mask)
        
        df = gribs.read(1)[0].data()[0]
        ev += df.data * (True ^ df.mask)
        
        df = gribs.read(1)[0].data()[0]
        p += df.data * (True ^ df.mask)
        
    gribs.close()
    return pe, ro, sro, ev, p


def get_era5_pe_etc2(era5_grib_file):
    '''
    Get 
        1:Potential evaporation:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        2:Runoff:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        3:Surface runoff:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        4:Evaporation:m of water equivalent (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
        5:Total precipitation:m (accum):regular_ll:surface:level 0:fcst time 23-24 hrs (accum):from 201912310000
    from file of path 'ERA5_eva_runoff_pre'.
    '''
    gribs = pygrib.open(era5_grib_file)
    messages = gribs.read()
    gribs.close()
    
    print('Reading data of hour %2d ...'%0)
    
    df = messages[0].data()[0]
    pe = df.data * (True ^ df.mask)
    
    df = messages[1].data()[0]
    ro = df.data * (True ^ df.mask)
    
    df = messages[2].data()[0]
    sro = df.data * (True ^ df.mask)
    
    df = messages[3].data()[0]
    ev = df.data * (True ^ df.mask)
    
    df = messages[4].data()[0]
    p = df.data * (True ^ df.mask)
    
    for hour in range(1, gribs.messages // 5):
        print('Reading data of hour %2d ...'%hour)
        
        df = messages[hour * 5 + 0].data()[0]
        pe += df.data * (True ^ df.mask)
    
        df = messages[hour * 5 + 1].data()[0]
        ro += df.data * (True ^ df.mask)
        
        df = messages[hour * 5 + 2].data()[0]
        sro += df.data * (True ^ df.mask)
        
        df = messages[hour * 5 + 3].data()[0]
        ev += df.data * (True ^ df.mask)
        
        df = messages[hour * 5 + 4].data()[0]
        p += df.data * (True ^ df.mask)
                
    return pe, ro, sro, ev, p


def tran_day_data_of_era5_pe_etc(era5_grib_file, dest_dump_file):
    if pathlib.Path(dest_dump_file).exists():
        print('Dest file exist, skip : %s'%dest_dump_file)
        return
    
    print('\nReading from file : %s'%era5_grib_file)
    print('-' * 40)
    pe, ro, sro, ev, p = get_era5_pe_etc(era5_grib_file)
    print('Dumping to   file : %s'%dest_dump_file)
    joblib.dump((pe, ro, sro, ev, p), dest_dump_file)
    print('-' * 10, 'Done', '-' * 10, '\n')
    return


#-----------------------------------------------------------------------------W


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print('\nUsage')
        print('-' * 40)
        print('python GribPrecipitationTran.py era5_grib_file_of_day dest_dump_file')
        print()
    else:
        era5_grib_file, dest_dump_file = sys.argv[1:]
        tran_day_data_of_era5_pe_etc(era5_grib_file, dest_dump_file)

    
'''
era5_grib_file = '/data/yule/ERA5_eva_runoff_pre/2020/01/eav_runoff_pre_2020_01_01.grib'
dump_file = '1.joblib'
show_messages(era5_grib_file)
'
tran_day_data_of_era5_t_e_vsm(era5_grib_file, dump_file)

import time

start = time.perf_counter()
k, e, u, d = get_era5_t_e_vsm(era5_grib_file)
print('Time used %s'%(time.perf_counter() - start))

start = time.perf_counter()
k, e, u, d = get_era5_t_e_vsm2(era5_grib_file)
print('Time used %s'%(time.perf_counter() - start))


#d = get_grib_data(era5_grib_file, 6)
#d = get_era5_t_e_vsm(grib_file)
#'''
